import pandas as pd
import numpy as np
import tushare as ts
from datetime import datetime, timedelta

# 设置tushare token
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()

# 获取上一个交易日
def get_last_trade_date():
    today = datetime.now()
    # 获取最近一个交易日
    trade_date = pro.trade_cal(exchange='', start_date=(today - timedelta(days=7)).strftime('%Y%m%d'),
                               end_date=today.strftime('%Y%m%d'), is_open='1')
    last_trade_date = trade_date.iloc[-1]['cal_date']
    return last_trade_date

last_trade_date = get_last_trade_date()
print(f"上一个交易日: {last_trade_date}")

# 获取所有股票的基本信息
stock_list = pro.stock_basic(exchange='', list_status='L')
stock_codes = stock_list['ts_code'].tolist()

# 获取上一个交易日的日线数据
def get_daily_data(trade_date):
    # 使用tushare的daily接口获取数据
    daily_data = pro.daily(trade_date=trade_date)
    return daily_data

daily_df = get_daily_data(last_trade_date)

# 合并股票基本信息
merged_df = pd.merge(daily_df, stock_list, on='ts_code', how='left')

# 数据概览
print("数据概览:")
print(f"股票数量: {len(merged_df)}")
print(f"数据时间: {last_trade_date}")
print("\n前5条数据:")
print(merged_df.head())

# 基本统计信息
print("\n基本统计信息:")
print(merged_df.describe())

# 各行业股票数量统计（如果数据中包含行业信息）
if 'industry' in merged_df.columns:
    industry_count = merged_df['industry'].value_counts()
    print("\n各行业股票数量:")
    print(industry_count)

# 涨跌幅分布分析
print("\n涨跌幅分布:")
print(merged_df['pct_chg'].describe())

# 成交量分析
print("\n成交量(手)分析:")
print(merged_df['vol'].describe())

# 成交额分析
print("\n成交额(千元)分析:")
print(merged_df['amount'].describe())

# 计算MA5和MA10
def calculate_ma(data):
    data['MA5'] = data['close'].rolling(window=5).mean()
    data['MA10'] = data['close'].rolling(window=10).mean()
    return data

# 计算RSI指标
def calculate_rsi(data, period=14):
    delta = data['close'].diff(1)
    gain = (delta.where(delta > 0, 0)).rolling(window=period).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=period).mean()
    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))
    return rsi

# 提取买点股票：5日线上穿10日线且RSI小于30
def find_buy_signals(data):
    data['Signal'] = (data['MA5'] > data['MA10']) & (data['RSI'] < 30)
    buy_signals = data[data['Signal']]
    return buy_signals

# 对每只股票计算MA5、MA10和RSI，并提取买点股票
buy_signals_list = []
all_stock_data = []

for ts_code in stock_codes[:10]:  # 示例：仅处理前10只股票（可根据需要调整）
    try:
        # 获取过去一段时间的日线数据（用于计算MA5、MA10和RSI）
        stock_data = pro.daily(ts_code=ts_code, start_date=(datetime.strptime(last_trade_date, '%Y%m%d') - timedelta(days=30)).strftime('%Y%m%d'), end_date=last_trade_date)
        stock_data = calculate_ma(stock_data)
        stock_data['RSI'] = calculate_rsi(stock_data)
        stock_data['ts_code'] = ts_code  # 添加股票代码列
        all_stock_data.append(stock_data)

        buy_signals = find_buy_signals(stock_data)
        if not buy_signals.empty:
            buy_signals_list.append(buy_signals)
    except Exception as e:
        print(f"处理股票 {ts_code} 时出错: {e}")

# 合并所有股票数据
all_stock_data_df = pd.concat(all_stock_data, ignore_index=True)

# 合并所有买点信号
if buy_signals_list:
    all_buy_signals = pd.concat(buy_signals_list, ignore_index=True)
    print("\n买点股票数据:")
    print(all_buy_signals)
else:
    print("\n未找到符合条件的买点股票。")

# 输出所有股票的MA5、MA10和RSI指标
print("\n所有股票的MA5、MA10和RSI指标:")
print(all_stock_data_df[['ts_code', 'trade_date', 'close', 'MA5', 'MA10', 'RSI']])